import tensorflow as tf
from PIL import ImageGrab
import numpy as np
import pyautogui
import time

# 加载训练好的模型
model = tf.keras.models.load_model('dino_gesture_model.keras')

# 图像参数
img_size = 100
img_channels = 1


def preprocess_image(image):

    # 转换为灰度图
    image = image.convert("L")
    # 调整大小
    image = image.resize((img_size, img_size))
    # 归一化
    image_np = np.array(image) / 255.0
    # 添加批次维度和通道维度
    image_np = np.expand_dims(image_np, axis=[0, -1])
    return image_np


def should_jump():

    screen = ImageGrab.grab(bbox=(280, 400, 650, 632)) 
    processed_screen = preprocess_image(screen)

    prediction = model.predict(processed_screen,verbose=0)
    predicted_class = np.argmax(prediction, axis=-1)[0]

    # 手动指定 jump 对应的类别索引
    class_indices = {'jump': 0, 'nojump': 1}
    jump_class = class_indices['jump']

    return predicted_class == jump_class


while True:
    if should_jump():
        pyautogui.press('space')  # 按下空格键跳跃
    time.sleep(0.01)  # 简单的延迟避免过度频繁的操作